{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "C4_W3_Lab_2_TFX_Custom_Components.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6x1ypzczQCwy"
      },
      "source": [
        "# Ungraded Lab: Developing Custom TFX Components\n",
        "\n",
        "[Tensorflow Extended (TFX)](https://www.tensorflow.org/tfx/) provides ready made components for typical steps in a machine learning workflow. Other courses in this specialization focus on specific components and in this lab, you will learn how to make your own. This will be useful in case your project has specific needs that fall outside the standard TFX components. It will make your pipelines more flexible while still leveraging the experiment tracking and orchestration that TFX provides. In particular, you will:\n",
        "\n",
        "* build a custom component using Python functions\n",
        "* build a custom component by reusing an existing TFX component\n",
        "* run a TFX pipeline locally using a built-in pipeline orchestrator\n",
        "\n",
        "To demonstrate, you will run the pipeline used in this [official tutorial](https://colab.research.google.com/github/tensorflow/tfx/blob/master/docs/tutorials/tfx/penguin_simple.ipynb) then modify it to have a custom component. Some of the discussions here are also taken from that tutorial to explain the motivation and point to additional resources.\n",
        "\n",
        "Let's begin!\n",
        "\n",
        "*Note: If you haven't taken other courses in this specialization and it's the first time you're using TFX, please see\n",
        "[Understanding TFX Pipelines](https://www.tensorflow.org/tfx/guide/understanding_tfx_pipelines)\n",
        "to get an overview of important concepts.*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fmgi8ZvQkScg"
      },
      "source": [
        "## Set Up\n",
        "We first need to install the TFX Python package and download\n",
        "the dataset which we will use for our model.\n",
        "\n",
        "### Upgrade Pip\n",
        "\n",
        "To avoid upgrading Pip in a system when running locally,\n",
        "check to make sure that we are running in Colab.\n",
        "Local systems can of course be upgraded separately."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "as4OTe2ukSqm"
      },
      "source": [
        "try:\n",
        "  import colab\n",
        "  !pip install --upgrade pip\n",
        "except:\n",
        "  pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MZOYTt1RW4TK"
      },
      "source": [
        "### Install TFX\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iyQtljP-qPHY"
      },
      "source": [
        "!pip install -U tfx"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EwT0nov5QO1M"
      },
      "source": [
        "**Note:** *Please do not proceed to the next steps without restarting the Runtime after installing `tfx`. You can do that by either pressing the `Restart Runtime` button at the end of the cell output above, or going to the `Runtime` button at the Colab toolbar above and selecting `Restart Runtime`.*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDnPgN8UJtzN"
      },
      "source": [
        "Check the TensorFlow and TFX versions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6jh7vKSRqPHb"
      },
      "source": [
        "import tensorflow as tf\n",
        "print('TensorFlow version: {}'.format(tf.__version__))\n",
        "from tfx import v1 as tfx\n",
        "print('TFX version: {}'.format(tfx.__version__))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aDtLdSkvqPHe"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "There are some variables used to define a pipeline. You can customize these\n",
        "variables as you want. By default all output from the pipeline will be\n",
        "generated under the current directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EcUseqJaE2XN"
      },
      "source": [
        "import os\n",
        "from absl import logging\n",
        "\n",
        "# Pipeline label\n",
        "PIPELINE_NAME = \"penguin-simple\"\n",
        "\n",
        "# Output directory to store artifacts generated from the pipeline.\n",
        "PIPELINE_ROOT = os.path.join('pipelines', PIPELINE_NAME)\n",
        "\n",
        "# Path to a SQLite DB file to use as an MLMD storage.\n",
        "METADATA_PATH = os.path.join('metadata', PIPELINE_NAME, 'metadata.db')\n",
        "\n",
        "# Output directory where created models from the pipeline will be exported.\n",
        "SERVING_MODEL_DIR = os.path.join('serving_model', PIPELINE_NAME)\n",
        "\n",
        "# Set default logging level.\n",
        "logging.set_verbosity(logging.INFO)  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8F2SRwRLSYGa"
      },
      "source": [
        "### Prepare example data\n",
        "\n",
        "The dataset you will use is the\n",
        "[Palmer Penguins dataset](https://allisonhorst.github.io/palmerpenguins/articles/intro.html).\n",
        "\n",
        "There are four numeric features in this dataset:\n",
        "\n",
        "- culmen_length_mm\n",
        "- culmen_depth_mm\n",
        "- flipper_length_mm\n",
        "- body_mass_g\n",
        "\n",
        "All features were already normalized to have range [0,1]. You will build a\n",
        "classification model which predicts the `species` of penguins. The code below creates a\n",
        "directory and copies the dataset to it. After running it, you should see the dataset in the Colab file explorer under the `data` folder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4fxMs6u86acP"
      },
      "source": [
        "import urllib.request\n",
        "\n",
        "# Create directory\n",
        "DATA_ROOT = 'data'\n",
        "!mkdir {DATA_ROOT}\n",
        "\n",
        "# Copy dataset to directory\n",
        "_data_url = 'https://raw.githubusercontent.com/tensorflow/tfx/master/tfx/examples/penguin/data/labelled/penguins_processed.csv'\n",
        "_data_filepath = os.path.join(DATA_ROOT, \"data.csv\")\n",
        "urllib.request.urlretrieve(_data_url, _data_filepath)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XlcD-gMhqlUv"
      },
      "source": [
        "## Running the pipeline using standard components\n",
        "\n",
        "The pipeline will consist of three essential TFX components and the graph will look like this:\n",
        "\n",
        "```\n",
        "ExampleGen -> Trainer -> Pusher\n",
        "``` \n",
        "\n",
        "The pipeline includes the most minimal ML workflow which is\n",
        "importing data (ExampleGen), training a model (Trainer) and exporting the trained model (Pusher)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SD3H-oc9rn1K"
      },
      "source": [
        "### Model training code\n",
        "\n",
        "You will first define the trainer module so the `Trainer` component can build the model and train it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EyinPEAKrm9z"
      },
      "source": [
        "_trainer_module_file = 'penguin_trainer.py'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jbQOc2VIqe-E"
      },
      "source": [
        "%%writefile {_trainer_module_file}\n",
        "\n",
        "from typing import List\n",
        "from absl import logging\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow_transform.tf_metadata import schema_utils\n",
        "\n",
        "from tfx import v1 as tfx\n",
        "from tfx_bsl.public import tfxio\n",
        "from tensorflow_metadata.proto.v0 import schema_pb2\n",
        "\n",
        "_FEATURE_KEYS = [\n",
        "    'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g'\n",
        "]\n",
        "_LABEL_KEY = 'species'\n",
        "\n",
        "_TRAIN_BATCH_SIZE = 20\n",
        "_EVAL_BATCH_SIZE = 10\n",
        "\n",
        "# Since we're not generating or creating a schema, we will instead create\n",
        "# a feature spec.  Since there are a fairly small number of features this is\n",
        "# manageable for this dataset.\n",
        "_FEATURE_SPEC = {\n",
        "    **{\n",
        "        feature: tf.io.FixedLenFeature(shape=[1], dtype=tf.float32)\n",
        "           for feature in _FEATURE_KEYS\n",
        "       },\n",
        "    _LABEL_KEY: tf.io.FixedLenFeature(shape=[1], dtype=tf.int64)\n",
        "}\n",
        "\n",
        "\n",
        "def _input_fn(file_pattern: List[str],\n",
        "              data_accessor: tfx.components.DataAccessor,\n",
        "              schema: schema_pb2.Schema,\n",
        "              batch_size: int = 200) -> tf.data.Dataset:\n",
        "  \"\"\"Generates features and label for training.\n",
        "\n",
        "  Args:\n",
        "    file_pattern: List of paths or patterns of input tfrecord files.\n",
        "    data_accessor: DataAccessor for converting input to RecordBatch.\n",
        "    schema: schema of the input data.\n",
        "    batch_size: representing the number of consecutive elements of returned\n",
        "      dataset to combine in a single batch\n",
        "\n",
        "  Returns:\n",
        "    A dataset that contains (features, indices) tuple where features is a\n",
        "      dictionary of Tensors, and indices is a single Tensor of label indices.\n",
        "  \"\"\"\n",
        "  return data_accessor.tf_dataset_factory(\n",
        "      file_pattern,\n",
        "      tfxio.TensorFlowDatasetOptions(\n",
        "          batch_size=batch_size, label_key=_LABEL_KEY),\n",
        "      schema=schema).repeat()\n",
        "\n",
        "\n",
        "def _build_keras_model() -> tf.keras.Model:\n",
        "  \"\"\"Creates a DNN Keras model for classifying penguin data.\n",
        "\n",
        "  Returns:\n",
        "    A Keras Model.\n",
        "  \"\"\"\n",
        "  # The model below is built with Functional API, please refer to\n",
        "  # https://www.tensorflow.org/guide/keras/overview for all API options.\n",
        "  inputs = [keras.layers.Input(shape=(1,), name=f) for f in _FEATURE_KEYS]\n",
        "  d = keras.layers.concatenate(inputs)\n",
        "  for _ in range(2):\n",
        "    d = keras.layers.Dense(8, activation='relu')(d)\n",
        "  outputs = keras.layers.Dense(3)(d)\n",
        "\n",
        "  model = keras.Model(inputs=inputs, outputs=outputs)\n",
        "  model.compile(\n",
        "      optimizer=keras.optimizers.Adam(1e-2),\n",
        "      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "      metrics=[keras.metrics.SparseCategoricalAccuracy()])\n",
        "\n",
        "  model.summary(print_fn=logging.info)\n",
        "  return model\n",
        "\n",
        "\n",
        "# TFX Trainer will call this function.\n",
        "def run_fn(fn_args: tfx.components.FnArgs):\n",
        "  \"\"\"Train the model based on given args.\n",
        "\n",
        "  Args:\n",
        "    fn_args: Holds args used to train the model as name/value pairs.\n",
        "  \"\"\"\n",
        "\n",
        "  # This schema is usually either an output of SchemaGen or a manually-curated\n",
        "  # version provided by pipeline author. A schema can also derived from TFT\n",
        "  # graph if a Transform component is used. In the case when either is missing,\n",
        "  # `schema_from_feature_spec` could be used to generate schema from very simple\n",
        "  # feature_spec, but the schema returned would be very primitive.\n",
        "  schema = schema_utils.schema_from_feature_spec(_FEATURE_SPEC)\n",
        "\n",
        "  train_dataset = _input_fn(\n",
        "      fn_args.train_files,\n",
        "      fn_args.data_accessor,\n",
        "      schema,\n",
        "      batch_size=_TRAIN_BATCH_SIZE)\n",
        "  eval_dataset = _input_fn(\n",
        "      fn_args.eval_files,\n",
        "      fn_args.data_accessor,\n",
        "      schema,\n",
        "      batch_size=_EVAL_BATCH_SIZE)\n",
        "\n",
        "  model = _build_keras_model()\n",
        "  model.fit(\n",
        "      train_dataset,\n",
        "      steps_per_epoch=fn_args.train_steps,\n",
        "      validation_data=eval_dataset,\n",
        "      validation_steps=fn_args.eval_steps)\n",
        "\n",
        "  # The result of the training should be saved in `fn_args.serving_model_dir`\n",
        "  # directory.\n",
        "  model.save(fn_args.serving_model_dir, save_format='tf')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1FiAttwdsXRn"
      },
      "source": [
        "### Write a pipeline definition\n",
        "\n",
        "Next, you will define a function to create a TFX pipeline. A [`Pipeline`](https://www.tensorflow.org/tfx/api_docs/python/tfx/v1/dsl/Pipeline) object represents a TFX pipeline which can be run using one of pipeline orchestration systems that TFX supports."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "78NRGIGNsNPD"
      },
      "source": [
        "def _create_pipeline(pipeline_name: str, pipeline_root: str, data_root: str,\n",
        "                     module_file: str, serving_model_dir: str,\n",
        "                     metadata_path: str) -> tfx.dsl.Pipeline:\n",
        "  \"\"\"Creates a three component penguin pipeline with TFX.\"\"\"\n",
        "  # Brings data into the pipeline.\n",
        "  example_gen = tfx.components.CsvExampleGen(input_base=data_root)\n",
        "\n",
        "  # Uses user-provided Python function that trains a model.\n",
        "  trainer = tfx.components.Trainer(\n",
        "      module_file=module_file,\n",
        "      examples=example_gen.outputs['examples'],\n",
        "      train_args=tfx.proto.TrainArgs(num_steps=100),\n",
        "      eval_args=tfx.proto.EvalArgs(num_steps=5))\n",
        "\n",
        "  # Pushes the model to a filesystem destination.\n",
        "  pusher = tfx.components.Pusher(\n",
        "      model=trainer.outputs['model'],\n",
        "      push_destination=tfx.proto.PushDestination(\n",
        "          filesystem=tfx.proto.PushDestination.Filesystem(\n",
        "              base_directory=serving_model_dir)))\n",
        "\n",
        "  # Following three components will be included in the pipeline.\n",
        "  components = [\n",
        "      example_gen,\n",
        "      trainer,\n",
        "      pusher,\n",
        "  ]\n",
        "\n",
        "  return tfx.dsl.Pipeline(\n",
        "      pipeline_name=pipeline_name,\n",
        "      pipeline_root=pipeline_root,\n",
        "      metadata_connection_config=tfx.orchestration.metadata\n",
        "      .sqlite_metadata_connection_config(metadata_path),\n",
        "      components=components)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k3A4E8YOsyv4"
      },
      "source": [
        "## Run the pipeline\n",
        "\n",
        "TFX supports multiple orchestrators to run pipelines.\n",
        "In this tutorial we will use `LocalDagRunner` which is included in the TFX\n",
        "Python package and runs pipelines on local environment.\n",
        "We often call TFX pipelines \"DAGs\" which stands for directed acyclic graph.\n",
        "\n",
        "`LocalDagRunner` provides fast iterations for developemnt and debugging.\n",
        "TFX also supports other orchestrators including Kubeflow Pipelines and Apache\n",
        "Airflow which are suitable for production use cases.\n",
        "\n",
        "See\n",
        "[TFX on Cloud AI Platform Pipelines](https://www.tensorflow.org/tfx/tutorials/tfx/cloud-ai-platform-pipelines)\n",
        "or\n",
        "[TFX Airflow Tutorial](https://www.tensorflow.org/tfx/tutorials/tfx/airflow_workshop)\n",
        "to learn more about other orchestration systems.\n",
        "\n",
        "The code below creates a `LocalDagRunner` and passes a `Pipeline` object created from the\n",
        "function you defined earlier. The pipeline runs directly and you can see logs for the progress of the pipeline including ML model training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S-TCt1Res8rH"
      },
      "source": [
        "tfx.orchestration.LocalDagRunner().run(\n",
        "  _create_pipeline(\n",
        "      pipeline_name=PIPELINE_NAME,\n",
        "      pipeline_root=PIPELINE_ROOT,\n",
        "      data_root=DATA_ROOT,\n",
        "      module_file=_trainer_module_file,\n",
        "      serving_model_dir=SERVING_MODEL_DIR,\n",
        "      metadata_path=METADATA_PATH))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YBkiw0wUtdBU"
      },
      "source": [
        "You should see \"INFO:absl:Component Pusher is finished.\" at the end of the\n",
        "logs if the pipeline finished successfully. Because `Pusher` component is the\n",
        "last component of the pipeline. You can also see that the `metadata` (metadata store), `pipelines` (pipeline root), and `serving_model` directories have been populated with the metadata, artifacts, and models that the pipeline generates.\n",
        "\n",
        "Now that you've ran this simple pipeline, you will see in the next sections how you can modify it to use custom components."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8WNr5JfHt-Sn"
      },
      "source": [
        "## Building Custom Components\n",
        "\n",
        "Let's say you want to modify the pipeline above to filter the data first before running the trainer. Without using a TFX component, you would run something like the code below. This will just get the rows where the `culmen_length_mm` feature is greater than `0.3`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k2EAJyXD3za_"
      },
      "source": [
        "import pandas as pd\n",
        "import glob\n",
        "\n",
        "# search directory for one or more CSVs\n",
        "files = glob.glob(f'{DATA_ROOT}/*.csv')\n",
        "\n",
        "# filter the dataset\n",
        "for file in files:\n",
        "  df = pd.read_csv(file, index_col=False)\n",
        "  filtered_df = df[df['culmen_length_mm'] > 0.3].reset_index(drop=True)\n",
        "\n",
        "# print latest modified file\n",
        "filtered_df"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cqr7pZ3lyYUd"
      },
      "source": [
        "You can save the dataset above to a different directory and point the TFX pipeline to it. That definitely works but you can include this code in a custom component as well so it's part of the TFX pipeline.\n",
        "\n",
        "As mentioned in the lectures, a TFX component has a driver, executor, and publisher. The driver and publisher interacts with the metadata store while the executor runs the actual processing. More often than not, you only need to modify the executor and that's what you'll be doing in the next sections."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8n6HkPOEzaSN"
      },
      "source": [
        "## Custom components through Python functions\n",
        "\n",
        "TFX provides a way to define an executor by just using a `component` decorator on your function. If you've done the Kubeflow Pipelines ungraded lab earlier this week, this will look very familiar. It uses the same concepts such as defining the inputs and outputs using type annotations then using those parameters in the function body.\n",
        "\n",
        "The component below basically uses the same filtering code you saw earlier but wraps it in a Python component. It defines two input parameters and outputs a dictionary with a String artifact (see function docstring for description). Because this component uses the same publisher as a standard TFX component, the output artifact will be saved in the metadata store as you'll see later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SwLgLLY3m51s"
      },
      "source": [
        "#import artifact type that you will use\n",
        "from tfx.types.standard_artifacts import String\n",
        "\n",
        "# import the decorator\n",
        "from tfx.dsl.component.experimental.decorators import component\n",
        "\n",
        "# import type annotations\n",
        "from tfx.dsl.component.experimental.annotations import InputArtifact, OutputArtifact, Parameter, OutputDict\n",
        "\n",
        "@component\n",
        "def CustomFilterComponent(input_base: Parameter[str], \n",
        "                    output_base: Parameter[str],\n",
        "                    ) -> OutputDict(output_path=str):\n",
        "  '''\n",
        "  Args:\n",
        "    input_base - location of the raw CSV\n",
        "    output_base - location where you want to save the filtered CSV\n",
        "  \n",
        "  Returns:\n",
        "    OutputDict:\n",
        "      output_path - String artifact that just holds the `output_base` value\n",
        "  '''\n",
        "  import pandas as pd\n",
        "  import glob\n",
        "  import os\n",
        "\n",
        "  # create the output base if it does not exist yet\n",
        "  if not os.path.exists(output_base):\n",
        "      os.mkdir(output_base)\n",
        "  \n",
        "  # search for CSVs in the input base\n",
        "  files = glob.glob(f'{input_base}/*.csv')\n",
        "\n",
        "  # loop through CSVs\n",
        "  for file in files:\n",
        "\n",
        "    # read the CSV\n",
        "    df = pd.read_csv(file, index_col=False)\n",
        "\n",
        "    # filter the data\n",
        "    filtered_df = df[df['culmen_length_mm'] > 0.3].reset_index(drop=True)\n",
        "\n",
        "    # compose output filename\n",
        "    filename = os.path.basename(file).replace('.csv','')\n",
        "    filtered_filename = f'{filename}_filtered.csv'\n",
        "  \n",
        "    # save filtered CSV to output base\n",
        "    filtered_df.to_csv(f'{output_base}/{filtered_filename}', index=False)\n",
        "\n",
        "  # define the output artifact\n",
        "  return {'output_path': output_base}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjKA_IWm29UA"
      },
      "source": [
        "You can now run your newly built component. You will just run a single node pipeline to prove that it works."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XGRzLRR-ukuN"
      },
      "source": [
        "# define a filter task\n",
        "filter_task = CustomFilterComponent(input_base=DATA_ROOT, \n",
        "                                        output_base=f'{DATA_ROOT}_filtered')\n",
        "\n",
        "# include the task\n",
        "components = [filter_task]\n",
        "\n",
        "# define a pipeline with only the single component\n",
        "pipeline = tfx.dsl.Pipeline(\n",
        "      pipeline_name=PIPELINE_NAME,\n",
        "      pipeline_root=PIPELINE_ROOT,\n",
        "      metadata_connection_config=tfx.orchestration.metadata\n",
        "      .sqlite_metadata_connection_config(METADATA_PATH),\n",
        "      components=components)\n",
        "\n",
        "# run the pipeline\n",
        "tfx.orchestration.LocalDagRunner().run(pipeline)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y2CxoUcu4Ne5"
      },
      "source": [
        "You should now see the filtered CSV in the `data_filtered` folder in the file explorer. As expected, you can also see that it has less lines than the original because of the filtering."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0cDdNBhku0j9"
      },
      "source": [
        "# number of rows in original csv\n",
        "!cat data/data.csv | wc -l\n",
        "\n",
        "# number of rows in filtered csv\n",
        "!cat data_filtered/data_filtered.csv | wc -l"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pUUTdRPU4xmZ"
      },
      "source": [
        "If you navigate to `pipelines/penguin-simple/CustomFilterComponent/output_path/`, you will also see a directory with the run id of the pipeline (e.g. `4`). If you click on the `value` file beneath it, you'll see the string value you saved which is just the `output_base` (i.e. `data_filtered`). You will want to feed this value to the `CsvExampleGen` component and that's what you'll do next."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JNkcASQ25Vgm"
      },
      "source": [
        "### Building a custom component using standard components\n",
        "\n",
        "If we try to use our new component with CsvExampleGen, **you will encounter an error** as shown below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L3WzxHyA56M8"
      },
      "source": [
        "# Define filter task\n",
        "filter_task = CustomFilterComponent(input_base=DATA_ROOT, \n",
        "                                        output_base=f'{DATA_ROOT}_filtered')\n",
        "\n",
        "# Try using the custom component with CsvExampleGen. This code will expectedly throw an error.\n",
        "try:\n",
        "  example_gen = tfx.components.CsvExampleGen(input_base=filter_task.outputs['output_path'])\n",
        "\n",
        "except Exception as e:\n",
        "  print(\"Error thrown as expected!\")\n",
        "  print(e)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OaZ9X-SH7QXG"
      },
      "source": [
        "As you can see, the output of our custom component cannot be accepted into `CsvExampleGen` because it is configured to accept a primitive string. TFX components generate and consume [`Channel`](https://www.tensorflow.org/tfx/api_docs/python/tfx/v1/dsl/Channel) objects and that's what our custom component outputs. For that, we need to build a custom data ingestion component that reuses the `CsvExampleGen` code but accepts a `Channel`. You will do that in the following sections."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K8U2wy1L76H7"
      },
      "source": [
        "### Standard ExampleGen code\n",
        "\n",
        "TFX is open source so the code for standard components can easily be found the public [repo](https://github.com/tensorflow/tfx/tree/89d3cb6c59acd0d487916bff703711815f1506b5/tfx). We placed links in the following sections of the actual files that you'll be modifying/overriding in case you want to compare what was changed for your custom ExampleGen. \n",
        "\n",
        "The class heirarchy for these components is pretty deep but in summary, you will only need to modify three:\n",
        "\n",
        "* Component Spec - this describes the inputs, outputs, and parameters passed on to the component\n",
        "* Executor - code for processing the inputs, outputs, and parameters\n",
        "* Component Class - puts everything together so your code can be run by the TFX runner."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ptdXl9mx8au2"
      },
      "source": [
        "### Modify the Component Spec\n",
        "\n",
        "First, you will need to modify the Component Spec. This file describes the parameters, inputs, and outputs of the standard components. Parameters are values supplied at runtime while inputs and outputs are values read from the metadata store. The original for ExampleGen is found [here](https://github.com/tensorflow/tfx/blob/89d3cb6c59acd0d487916bff703711815f1506b5/tfx/types/standard_component_specs.py#L187).\n",
        "\n",
        "You will need to implement the same for our custom ExampleGen but it should accept a `Channel` parameter The revised version is shown below. Notice that we revised the `INPUTS` dictionary to have a `ChannelParameter` whereas in the original, all are just in the `PARAMETERS` dictionary."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aZlH_YRp0mVY"
      },
      "source": [
        "from tfx.types.component_spec import ChannelParameter\n",
        "from tfx.types.component_spec import ExecutionParameter\n",
        "from tfx.types.component_spec import ComponentSpec\n",
        "from tfx.types import standard_artifacts\n",
        "from tfx.proto import example_gen_pb2\n",
        "from tfx.proto import range_config_pb2\n",
        "\n",
        "\n",
        "# Key for example_gen input that we want to use\n",
        "INPUT_BASE_KEY = 'input_base'\n",
        "\n",
        "# Other keys\n",
        "INPUT_CONFIG_KEY = 'input_config'\n",
        "OUTPUT_CONFIG_KEY = 'output_config'\n",
        "OUTPUT_DATA_FORMAT_KEY = 'output_data_format'\n",
        "RANGE_CONFIG_KEY = 'range_config'\n",
        "CUSTOM_CONFIG_KEY = 'custom_config'\n",
        "EXAMPLES_KEY = 'examples'\n",
        "\n",
        "class MyCustomExampleGenSpec(ComponentSpec):\n",
        "  \"\"\"File-based ExampleGen component spec.\"\"\"\n",
        "  \n",
        "  PARAMETERS = {\n",
        "      INPUT_CONFIG_KEY:\n",
        "          ExecutionParameter(type=example_gen_pb2.Input),\n",
        "      OUTPUT_CONFIG_KEY:\n",
        "          ExecutionParameter(type=example_gen_pb2.Output),\n",
        "      OUTPUT_DATA_FORMAT_KEY:\n",
        "          ExecutionParameter(type=int),\n",
        "      CUSTOM_CONFIG_KEY:\n",
        "          ExecutionParameter(type=example_gen_pb2.CustomConfig, optional=True),\n",
        "      RANGE_CONFIG_KEY:\n",
        "          ExecutionParameter(type=range_config_pb2.RangeConfig, optional=True),\n",
        "  }\n",
        "\n",
        "  # Now accepts a channel\n",
        "  INPUTS = {\n",
        "      INPUT_BASE_KEY:\n",
        "          ChannelParameter(type=standard_artifacts.String),\n",
        "  }\n",
        "  \n",
        "  OUTPUTS = {\n",
        "      EXAMPLES_KEY: ChannelParameter(type=standard_artifacts.Examples),\n",
        "  }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v_YXOyJc-3Tu"
      },
      "source": [
        "### Customize the Executor\n",
        "\n",
        "With that, you should now modify the executor code to take note of this change of input types. Instead of looking at just the parameters, it should also look into Channel inputs passed onto the component.\n",
        "\n",
        "Executor classes are executed by TFX starting with the `Do` function so you will need to modify that. The original Executor for `CsvExampleGen` can be found [here](https://github.com/tensorflow/tfx/blob/89d3cb6c59acd0d487916bff703711815f1506b5/tfx/components/example_gen/csv_example_gen/executor.py) and it inherits the base class [here](https://github.com/tensorflow/tfx/blob/89d3cb6c59acd0d487916bff703711815f1506b5/tfx/components/example_gen/base_example_gen_executor.py#L132). The base class includes the `Do()` function and that's what you'll be overriding in your new custom executor below. \n",
        "\n",
        "Basically, you're using all the functions defined in the standard component but you're modifying it so it can find the `input_base` value from the Channel inputs.\n",
        "\n",
        "*Take note that this tutorial prioritizes code brevity. In your projects, you may take a different approach such as modifying the [`_CsvToExample`](https://github.com/tensorflow/tfx/blob/89d3cb6c59acd0d487916bff703711815f1506b5/tfx/components/example_gen/csv_example_gen/executor.py#L124) code to look for the string value in the `input_dict` instead of `exec_properties`. Doing that here in this Colab will result in very long code blocks so the shorter approach is taken.*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c_agzAOu42v_"
      },
      "source": [
        "from typing import Any, Dict, Iterable, List, Text, Union\n",
        "\n",
        "from tfx.components.example_gen.csv_example_gen.executor import Executor as CsvExampleGenExecutor\n",
        "from tfx.types import standard_component_specs\n",
        "from tfx.types import artifact_utils\n",
        "from tfx import types\n",
        "\n",
        "class MyCustomExecutor(CsvExampleGenExecutor):\n",
        "  \"\"\"Generic TFX CSV example gen executor.\"\"\"\n",
        "\n",
        "  def Do(\n",
        "      self,\n",
        "      input_dict: Dict[Text, List[types.Artifact]],\n",
        "      output_dict: Dict[Text, List[types.Artifact]],\n",
        "      exec_properties: Dict[Text, Any],\n",
        "  ) -> None:\n",
        "    \"\"\"Take input data source and generates serialized data splits.\n",
        "    The output is intended to be serialized tf.train.Examples or\n",
        "    tf.train.SequenceExamples protocol buffer in gzipped TFRecord format,\n",
        "    but subclasses can choose to override to write to any serialized records\n",
        "    payload into gzipped TFRecord as specified, so long as downstream\n",
        "    component can consume it. The format of payload is added to\n",
        "    `payload_format` custom property of the output Example artifact.\n",
        "    Args:\n",
        "      input_dict: Input dict from input key to a list of Artifacts. Depends on\n",
        "        detailed example gen implementation.\n",
        "        - input_base: an external directory containing the data files.\n",
        "      output_dict: Output dict from output key to a list of Artifacts.\n",
        "        - examples: splits of serialized records.\n",
        "      exec_properties: A dict of execution properties. Depends on detailed\n",
        "        example gen implementation.\n",
        "        - input_config: JSON string of example_gen_pb2.Input instance,\n",
        "          providing input configuration.\n",
        "        - output_config: JSON string of example_gen_pb2.Output instance,\n",
        "          providing output configuration.\n",
        "        - output_data_format: Payload format of generated data in output\n",
        "          artifact, one of example_gen_pb2.PayloadFormat enum.\n",
        "    Returns:\n",
        "      None\n",
        "    \"\"\"\n",
        "    self._log_startup(input_dict, output_dict, exec_properties)\n",
        "\n",
        "    # Get the artifact from the Channel input\n",
        "    filter_component_artifact = artifact_utils.get_single_instance(\n",
        "        input_dict[standard_component_specs.INPUT_BASE_KEY])\n",
        "    \n",
        "    # Put the input string value into the exec_properties fictionary\n",
        "    exec_properties[standard_component_specs.INPUT_BASE_KEY] = filter_component_artifact.value\n",
        "    \n",
        "    # execute superclass\n",
        "    super(MyCustomExecutor, self).Do(input_dict=input_dict, output_dict=output_dict, exec_properties=exec_properties)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "inMjz8qDAJyz"
      },
      "source": [
        "### Define the Component class\n",
        "\n",
        "Lastly, you will need to put everything together in a class. This will be the one you instantiate so you can run the component later. For comparison, the original `CsvExampleGen` component class is found [here](https://github.com/tensorflow/tfx/blob/89d3cb6c59acd0d487916bff703711815f1506b5/tfx/components/example_gen/csv_example_gen/component.py) and it inherits the `FileBasedExampleGen` from [here](https://github.com/tensorflow/tfx/blob/89d3cb6c59acd0d487916bff703711815f1506b5/tfx/components/example_gen/component.py#L115). The revised version is shown below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wOlYnQumyWfX"
      },
      "source": [
        "from typing import Any, Dict, Optional, Text, Union\n",
        "\n",
        "from tfx.dsl.components.base import base_beam_component\n",
        "from tfx.dsl.components.base import executor_spec\n",
        "\n",
        "from tfx.proto import example_gen_pb2\n",
        "from tfx.proto import range_config_pb2\n",
        "\n",
        "from tfx import types\n",
        "from tfx.components.example_gen import utils\n",
        "\n",
        "class MyCustomExampleGen(base_beam_component.BaseBeamComponent):\n",
        "\n",
        "  # Define the Spec class and executor spec using the functions and\n",
        "  # classes you defined earlier.\n",
        "  SPEC_CLASS = MyCustomExampleGenSpec\n",
        "  EXECUTOR_SPEC = executor_spec.BeamExecutorSpec(MyCustomExecutor)\n",
        "\n",
        "  # Define init function. Notice that `input_base` now accepts a Channel.\n",
        "  def __init__(self,\n",
        "               input_base: types.Channel = None,\n",
        "               input_config: Optional[Union[example_gen_pb2.Input,\n",
        "                                            Dict[Text, Any]]] = None,\n",
        "               output_config: Optional[Union[example_gen_pb2.Output,\n",
        "                                             Dict[Text, Any]]] = None,\n",
        "               range_config: Optional[Union[range_config_pb2.RangeConfig,\n",
        "                                            Dict[Text, Any]]] = None,\n",
        "               output_data_format: Optional[int] = example_gen_pb2.FORMAT_TF_EXAMPLE):\n",
        "    \"\"\"Customized ExampleGen component.\n",
        "    Args:\n",
        "      input_base: an external directory containing the CSV files. Accepts a Channel\n",
        "        from a previous TFX component.\n",
        "      input_config: An example_gen_pb2.Input instance, providing input\n",
        "        configuration. If unset, the files under input_base will be treated as a\n",
        "        single split. If any field is provided as a RuntimeParameter,\n",
        "        input_config should be constructed as a dict with the same field names\n",
        "        as Input proto message.\n",
        "      output_config: An example_gen_pb2.Output instance, providing output\n",
        "        configuration. If unset, default splits will be 'train' and 'eval' with\n",
        "        size 2:1. If any field is provided as a RuntimeParameter, output_config\n",
        "        should be constructed as a dict with the same field names as Output\n",
        "        proto message.\n",
        "      range_config: An optional range_config_pb2.RangeConfig instance,\n",
        "        specifying the range of span values to consider. If unset, driver will\n",
        "        default to searching for latest span with no restrictions.\n",
        "    \"\"\"\n",
        "    \n",
        "    # Configure inputs and outputs.\n",
        "    input_config = input_config or utils.make_default_input_config()\n",
        "    output_config = output_config or utils.make_default_output_config(\n",
        "        input_config)\n",
        "    \n",
        "    # Define output type.\n",
        "    example_artifacts = types.Channel(type=standard_artifacts.Examples)\n",
        "    \n",
        "    # Pass input arguments to your custom ExampleGen spec.\n",
        "    spec = MyCustomExampleGenSpec(\n",
        "        input_base=input_base,\n",
        "        input_config=input_config,\n",
        "        output_config=output_config,\n",
        "        range_config=range_config,\n",
        "        output_data_format=output_data_format,\n",
        "        examples=example_artifacts)\n",
        "    \n",
        "    # This will check if the values passed are the correct type else\n",
        "    # it will throw the error you saw earlier.\n",
        "    super(MyCustomExampleGen, self).__init__(\n",
        "        spec=spec)\n",
        "  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n2oFBiFZggLS"
      },
      "source": [
        "You can now use the custom component (`MyCustomExampleGen`) in your code as shown below. It will no longer get an error because you reconfigured the `input_base` to accept a channel."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DXOlqQ_9y-rL"
      },
      "source": [
        "# Filter the dataset\n",
        "filter_task = CustomFilterComponent(input_base=DATA_ROOT, \n",
        "                                        output_base=f'{DATA_ROOT}_filtered')\n",
        "\n",
        "# Use the output of filter_task to know the input_base for this custom ExampleGen\n",
        "custom_example_gen_task = MyCustomExampleGen(input_base=filter_task.outputs['output_path'])\n",
        "\n",
        "# Define components to include\n",
        "components = [filter_task,\n",
        "              custom_example_gen_task]\n",
        "\n",
        "# Create the pipeline\n",
        "pipeline = tfx.dsl.Pipeline(\n",
        "      pipeline_name=PIPELINE_NAME,\n",
        "      pipeline_root=PIPELINE_ROOT,\n",
        "      metadata_connection_config=tfx.orchestration.metadata\n",
        "      .sqlite_metadata_connection_config(METADATA_PATH),\n",
        "      components=components)\n",
        "\n",
        "# Run the pipeline\n",
        "tfx.orchestration.LocalDagRunner().run(pipeline)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y8nfx3annFTC"
      },
      "source": [
        "As a sanity check, you can compute the number of examples for both the training and eval splits. It should equal the number of examples found in your filtered CSV. You can use the code below by replacing the `EXECUTION_ID` with the ID shown in your latest run. You can see it in the last three lines of the output cell above. For example:\n",
        "\n",
        "```\n",
        ")]}) for execution 16  # ---> **16 is the EXECUTION ID**\n",
        "INFO:absl:MetadataStore with DB connection initialized\n",
        "INFO:absl:Component MyCustomExampleGen is finished.\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6DxCauSIPl_X"
      },
      "source": [
        "EXECUTION_ID = 16 # PLACE THE EXECUTION ID HERE\n",
        "\n",
        "# Create a `TFRecordDataset` to read these files\n",
        "train_dataset = tf.data.TFRecordDataset(f'{PIPELINE_ROOT}/MyCustomExampleGen/examples/{EXECUTION_ID}/Split-train/data_tfrecord-00000-of-00001.gz', compression_type=\"GZIP\")\n",
        "eval_dataset = tf.data.TFRecordDataset(f'{PIPELINE_ROOT}/MyCustomExampleGen/examples/{EXECUTION_ID}/Split-eval/data_tfrecord-00000-of-00001.gz', compression_type=\"GZIP\")\n",
        "\n",
        "# Get number of records for each dataset (only use for small datasets to avoid memory issues)\n",
        "num_train_data = len(list(train_dataset))\n",
        "num_eval_data = len(list(eval_dataset))\n",
        "\n",
        "# Get the total\n",
        "total_examples = num_train_data + num_eval_data\n",
        "\n",
        "print(f'total number of examples: {total_examples}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V8jE62T0ry43"
      },
      "source": [
        "## Wrap Up\n",
        "\n",
        "In this lab, you were able to use custom components to create a pipeline. This shows that you can go outside the standard TFX components if your project calls for it. To know more about custom components, you can read more [here](https://www.tensorflow.org/tfx/guide/understanding_custom_components) and see the examples [here](https://github.com/tensorflow/tfx/tree/2e41786328f5b69720e90ec4d9ecae500f5c157a/tfx/examples/custom_components)."
      ]
    }
  ]
}